import os
import yaml
import time
import argparse
import random
import numpy as np
import torch
import matplotlib.pyplot as plt
from typing import Dict, List, Tuple, Any

from env.grid_env import GridMap
from agent.dqn_agent import DQNAgent
from planner.dqn_planner import DQNPlanner
from cbs.cbs_solver import CBSSolver
from tasks.task_generator import TaskGenerator
from simulation.simulator import Simulator
from evaluation.metrics import Metrics
from utils.logger import get_logger
from agent.train_simple import create_warehouse_map  # 导入仓储地图创建函数


def load_config(config_path: str) -> Dict[str, Any]:
    """
    加载配置文件
    
    Args:
        config_path: 配置文件路径
        
    Returns:
        Dict[str, Any]: 配置参数
    """
    with open(config_path, 'r', encoding='utf-8') as f:
        config = yaml.safe_load(f)
    return config


def set_seed(seed: int):
    """
    设置随机种子
    
    Args:
        seed: 随机种子
    """
    if seed is None:
        return
        
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)


def create_grid_map(config: Dict[str, Any], seed: int = None, use_warehouse: bool = False) -> GridMap:
    """
    创建网格地图
    
    Args:
        config: 配置参数
        seed: 随机种子
        use_warehouse: 是否使用仓储地图
        
    Returns:
        GridMap: 网格地图
    """
    width = config['map']['width']
    height = config['map']['height']
    obstacle_ratio = config['map']['obstacle_ratio']
    
    if use_warehouse:
        # 使用仓储地图
        grid_map, _ = create_warehouse_map(width=width, height=height, seed=seed)
        print(f"创建仓储网格地图: {width}x{height}")
        return grid_map
    else:
        # 使用普通地图
        return GridMap(width=width, height=height, obstacle_ratio=obstacle_ratio, seed=seed)


def initialize_robots(grid_map: GridMap, config: Dict[str, Any]) -> Dict[int, Tuple[int, int]]:
    """
    初始化机器人位置
    
    Args:
        grid_map: 网格地图
        config: 配置参数
        
    Returns:
        Dict[int, Tuple[int, int]]: 机器人位置字典 {robot_id: (x, y)}
    """
    robot_count = config['robots']['count']
    initial_positions = config['robots'].get('initial_positions', [])
    
    robot_positions = {}
    
    # 如果提供了初始位置，使用它们
    if initial_positions and len(initial_positions) >= robot_count:
        for i in range(robot_count):
            pos = tuple(initial_positions[i])
            if grid_map.is_valid_position(pos):
                robot_positions[i] = pos
    
    # 如果初始位置不足或无效，随机生成
    while len(robot_positions) < robot_count:
        x = random.randint(0, grid_map.width - 1)
        y = random.randint(0, grid_map.height - 1)
        pos = (x, y)
        
        if grid_map.is_valid_position(pos) and pos not in robot_positions.values():
            robot_positions[len(robot_positions)] = pos
    
    return robot_positions


def run_experiment(config: Dict[str, Any], model_path: str, seed: int = None, use_warehouse: bool = True):
    """
    运行一次实验
    
    Args:
        config: 配置参数
        model_path: DQN模型路径
        seed: 随机种子
        use_warehouse: 是否使用仓储地图
    """
    logger = get_logger()
    logger.info("开始实验...")
    
    # 设置随机种子
    if seed is not None:
        set_seed(seed)
        logger.info(f"设置随机种子: {seed}")
    
    # 创建网格地图 - 使用仓储地图
    grid_map = create_grid_map(config, seed, use_warehouse=use_warehouse)
    logger.info(f"创建网格地图: {grid_map.width}x{grid_map.height}, 障碍比例: {grid_map.obstacle_ratio}")
    
    # 初始化机器人
    robot_positions = initialize_robots(grid_map, config)
    logger.info(f"初始化 {len(robot_positions)} 个机器人")
    
    # 创建任务生成器
    task_generator = TaskGenerator(grid_map, config['tasks'])
    
    # 创建CBS求解器
    cbs_solver = CBSSolver(grid_map, model_path, config['cbs'])
    logger.info(f"创建CBS求解器，最大迭代次数: {cbs_solver.max_iterations}, 超时: {cbs_solver.timeout}秒")
    
    # 创建仿真器
    simulator = Simulator(grid_map, config['simulation'])
    
    # 创建评估指标
    metrics = Metrics()
    
    # 运行不同数量的任务
    task_counts = [1, 2, 3, len(robot_positions)]
    
    # 记录每个任务的结果
    all_results = []
    
    for task_count in task_counts:
        logger.info(f"\n===== 运行任务数量: {task_count} =====")
        
        # 生成任务
        tasks = task_generator.generate_balanced_tasks(task_count, robot_positions)
        logger.info(f"生成 {len(tasks)} 个任务")
        
        # 记录起始时间
        start_time = time.time()
        
        # 使用CBS求解
        solution = cbs_solver.find_solution(tasks)
        
        # 记录调度时间
        scheduling_time = time.time() - start_time
        logger.info(f"CBS求解完成，耗时: {scheduling_time:.4f}秒")
        
        # 如果没有解决方案，跳过
        if not solution:
            logger.warning("CBS无法找到解决方案，跳过此任务")
            continue
        
        # 提取目标位置
        goals = {agent_id: task[1] for agent_id, task in tasks.items()}
        
        # 加载解决方案到仿真器
        simulator.load_solution(solution, goals)
        
        # 记录执行开始时间
        exec_start_time = time.time()
        
        # 运行仿真
        simulator.run()
        
        # 记录执行时间
        execution_time = time.time() - exec_start_time
        
        # 获取统计信息
        stats = simulator.get_statistics()
        
        # 计算路径长度
        path_lengths = {agent_id: len(path) for agent_id, path in solution.items()}
        
        # 计算任务完成率
        success_rate = (stats['completed_robots'] / stats['total_robots']) * 100
        
        # 记录评估指标
        metrics.add_result(
            path_lengths=path_lengths,
            success_rate=success_rate,
            conflict_count=stats['conflicts'],
            scheduling_time=scheduling_time,
            execution_time=execution_time,
            task_count=task_count
        )
        
        # 记录当前任务的结果
        all_results.append({
            'task_count': task_count,
            'success_rate': success_rate,
            'conflicts': stats['conflicts'],
            'makespan': stats['makespan'],
            'sum_of_costs': stats['sum_of_costs'],
            'completed_robots': stats['completed_robots'],
            'total_robots': stats['total_robots']
        })
        
        # 打印统计信息
        logger.info(f"完成任务数: {stats['completed_robots']}/{stats['total_robots']}")
        logger.info(f"成功率: {success_rate:.2f}%")
        logger.info(f"冲突数: {stats['conflicts']}")
        logger.info(f"完工时间: {stats['makespan']}")
        logger.info(f"总代价: {stats['sum_of_costs']}")
    
    # 打印评估指标
    metrics.print_metrics()
    
    # 绘制评估指标图表
    metrics.plot_metrics(save_path="results/metrics_plot.png")
    
    # 保存评估指标
    metrics.save_metrics("results/metrics.txt")
    
    # 创建结果总结图表
    create_summary_chart(all_results)
    
    logger.info("实验完成")


def create_summary_chart(results: List[Dict[str, Any]]):
    """
    创建实验结果总结图表
    
    Args:
        results: 实验结果列表
    """
    if not results:
        return
    
    # 创建图表
    fig, axs = plt.subplots(2, 2, figsize=(14, 10))
    
    # 提取数据
    task_counts = [r['task_count'] for r in results]
    success_rates = [r['success_rate'] for r in results]
    conflicts = [r['conflicts'] for r in results]
    makespans = [r['makespan'] for r in results]
    sum_of_costs = [r['sum_of_costs'] for r in results]
    
    # 绘制任务成功率
    axs[0, 0].bar(task_counts, success_rates, color='green')
    axs[0, 0].set_title('任务成功率 (%)')
    axs[0, 0].set_xlabel('任务数量')
    axs[0, 0].set_ylabel('成功率 (%)')
    axs[0, 0].grid(True, linestyle='--', alpha=0.7)
    
    # 绘制冲突数
    axs[0, 1].bar(task_counts, conflicts, color='red')
    axs[0, 1].set_title('冲突数')
    axs[0, 1].set_xlabel('任务数量')
    axs[0, 1].set_ylabel('冲突数')
    axs[0, 1].grid(True, linestyle='--', alpha=0.7)
    
    # 绘制完工时间
    axs[1, 0].bar(task_counts, makespans, color='blue')
    axs[1, 0].set_title('完工时间')
    axs[1, 0].set_xlabel('任务数量')
    axs[1, 0].set_ylabel('时间步')
    axs[1, 0].grid(True, linestyle='--', alpha=0.7)
    
    # 绘制总代价
    axs[1, 1].bar(task_counts, sum_of_costs, color='purple')
    axs[1, 1].set_title('总代价')
    axs[1, 1].set_xlabel('任务数量')
    axs[1, 1].set_ylabel('代价')
    axs[1, 1].grid(True, linestyle='--', alpha=0.7)
    
    # 设置标题
    fig.suptitle('多机器人路径规划实验结果总结', fontsize=16)
    
    # 调整布局
    plt.tight_layout()
    plt.subplots_adjust(top=0.9)
    
    # 保存图表
    plt.savefig("results/experiment_summary.png", dpi=300)
    print(f"已保存实验结果总结图表到 results/experiment_summary.png")


def main():
    parser = argparse.ArgumentParser(description="CBS+DQN多机器人路径规划系统")
    parser.add_argument('--config', type=str, default='config/map_config.yaml', help='配置文件路径')
    parser.add_argument('--model', type=str, default='models/dqn_model_final.pth', help='DQN模型路径')
    parser.add_argument('--seed', type=int, default=None, help='随机种子')
    parser.add_argument('--train', action='store_true', help='是否训练DQN模型')
    parser.add_argument('--warehouse', action='store_true', default=True, help='是否使用仓储地图')
    args = parser.parse_args()
    
    # 创建结果目录
    os.makedirs("results", exist_ok=True)
    os.makedirs("logs", exist_ok=True)
    
    # 加载配置
    config = load_config(args.config)
    
    if args.train:
        # 训练仓储环境下的DQN模型
        from agent.train_simple import train_simple_dqn
        
        print("开始训练小规模网格环境DQN模型...")
        # 使用小规模网格 10x10 而不是配置文件中的大小
        width = 10
        height = 10
        # 修改障碍物比例为0.1，在小网格中避免过于拥挤
        train_simple_dqn(width=width, height=height, episodes=200, seed=args.seed, use_warehouse=False, max_steps=50)
    else:
        # 运行实验
        if not os.path.exists(args.model):
            print(f"错误: 模型文件 {args.model} 不存在!")
            print("请先使用 --train 参数训练模型，或者指定正确的模型路径。")
            return
        
        run_experiment(config, args.model, seed=args.seed, use_warehouse=args.warehouse)


if __name__ == "__main__":
    main() 